#include <clusnet/net/TcpClient.h>
#include <clusnet/base/Logging.h>
#include <list>
#include <memory>
#include <mutex>
#include <iostream>

class stats final 
{
    public:
        stats()
            : mutex_(),
              totalBytesWritten_(0),
              totalBytesRead_(0)
        {}
        ~stats() = default;
        stats(const stats&) = delete;
        stats& operator=(const stats&) = delete;

        void add(size_t bytesWritten, size_t bytesRead)
        {
            std::lock_guard<std::mutex> lock(mutex_);
            totalBytesWritten_ += bytesWritten;
            totalBytesRead_ += bytesRead;
        }

        void print(int spendTime)
        {
            std::lock_guard<std::mutex> lock(mutex_);
            LOG(INFO) << totalBytesWritten_ << " total bytes written, " << totalBytesRead_ << " total bytes read, " << static_cast<double>(totalBytesRead_) / (spendTime * 1024 * 1024) << " MiB/s throughput";
        }

    private:
        std::mutex  mutex_;
        size_t      totalBytesWritten_;
        size_t      totalBytesRead_;
};

class session final
{
    public:
        session(boost::asio::io_context& io, size_t blockSize, stats& stat, const char* host, const char* port)
            : bytesWritten_(0),
              bytesRead_(0),
              client_(new clusnet::net::TcpClient(io, host, port)),
              stat_(stat)
        {
            client_->setNewConnectionCallback(std::bind(&session::handleConnect, this, std::placeholders::_1));
            client_->setCloseCallback(std::bind(&session::handleClose, this, std::placeholders::_1));
            client_->setWriteCompleteCallback(std::bind(&session::handleSend, this, std::placeholders::_1, std::placeholders::_2));
            client_->setMessageCallback(std::bind(&session::handleRead, this, std::placeholders::_1, std::placeholders::_2));

            std::string data(blockSize, 'a');            
            clusnet::net::TcpDataPacketStream packetStream;
            packetStream.append(data.c_str(), data.length());
            data_ = std::make_shared<clusnet::net::TcpDataPacket>(std::move(packetStream).data());
        }
        ~session() = default;
        session(const session&) = delete;
        session& operator= (const session&) = delete;

        void start()
        {
            client_->safeStartConnect();
        }

        void close()
        {
            client_->safeClose();
        }

    private:
        void handleConnect(clusnet::net::TcpConnectionPtr& conn)
        {
            conn->safeStartReceive();
            client_->safeAsyncWrite(data_);
        }

        void handleClose(clusnet::net::TcpConnectionPtr& conn)
        {
            stat_.add(bytesWritten_, bytesRead_);
        }

        void handleRead(clusnet::net::TcpConnectionPtr& conn, clusnet::net::TcpUniqDataPacketPtr&& data)
        {
            bytesRead_ += data->dataLength();
            client_->safeAsyncWrite(std::move(data));
        }

        void handleSend(clusnet::net::TcpConnectionPtr&, const clusnet::net::TcpDataPacketPtr& data)
        {
            bytesWritten_ += data->dataLength();
        }

    private:
        size_t  bytesWritten_;
        size_t  bytesRead_;
        clusnet::net::TcpClientPtr      client_;
        stats&  stat_;
        clusnet::net::TcpDataPacketPtr  data_;
};

class client final 
{
    public:
        client(size_t threadCount, size_t sessionCount, int timeout, size_t blockSize, const char* host, const char* port)
            : ioContexts_(threadCount),
              threadCount_(threadCount),
              stopTimer_(ioContexts_[0]),
              timeout_(timeout)
        {  
            for (size_t i = 0; i < sessionCount; ++i)
            {
                auto& io = ioContexts_[i % threadCount_];
                std::unique_ptr<session> new_session(new session(io, blockSize, stat_, host, port));
                sessions_.emplace_back(std::move(new_session));
            }
            threads_.reserve(threadCount_);
        }

        ~client()
        {
            sessions_.clear(); 
        }

        void start()
        {
            for (auto& sess : sessions_)
            {
                sess->start();
            }
            LOG(INFO) << "sessions start";

            for (size_t i = 0; i < threadCount_; ++i)
            {
                threads_.emplace_back(new std::thread([this, i]
                {
                    auto& io = ioContexts_[i];
                    io.run();
                }));
            }
            LOG(INFO) << "thread start";
        }

        void wait()
        {
            stopTimer_.expires_after(std::chrono::seconds(timeout_));
            stopTimer_.wait();
            stop();
        } 

    private:
        void stop()
        {
            for (auto& sess : sessions_)
            {
                sess->close();
            }
            std::for_each(threads_.begin(), threads_.end(), [](std::unique_ptr<std::thread>& t) { t->join(); });
            stat_.print(timeout_);
        }

    private:
        std::vector<boost::asio::io_context>    ioContexts_;
        std::list<std::unique_ptr<session>> sessions_;  // 需要在iocontex销毁前销毁所有的session
        std::vector<std::unique_ptr<std::thread>>            threads_;
        size_t  threadCount_;
        boost::asio::steady_timer stopTimer_;
        int timeout_;
        stats   stat_;
};

int main(int argc, char* argv[])
{
    InitLogging("PingPongClient", ".");
    if (argc < 7)
    {
        std::cerr << "Usage: client <host> <port> <threadNum> <sessionNum> <timeout> <blockSize>" << std::endl;
        return 1;
    }
    const char* host = argv[1];
    const char* port = argv[2];
    size_t threadNum = std::atoi(argv[3]);
    size_t sessionNum = std::atoi(argv[4]);
    size_t timeout = std::atoi(argv[5]);
    size_t blockSize = std::atoi(argv[6]);

    client  c(threadNum, sessionNum, timeout, blockSize, host, port);
    c.start();
    c.wait();

    return 0;
}
